import torch
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Rectangle
from matplotlib.font_manager import FontProperties

# 加载中文字体（避免中文乱码）
font = FontProperties(fname="/usr/share/fonts/truetype/wqy/wqy-microhei.ttc", size=14)  # Ubuntu系统
# font = FontProperties(fname="C:/Windows/Fonts/simhei.ttf", size=14)  # Windows系统


# 1. 棋盘常量定义
BOARD_ROWS = 10  # 中国象棋10行
BOARD_COLS = 9   # 9列
CELL_SIZE = 60   # 每个格子的像素大小
WINDOW_WIDTH = BOARD_COLS * CELL_SIZE
WINDOW_HEIGHT = BOARD_ROWS * CELL_SIZE

# 棋子编码（正=红方，负=黑方）
PIECE_CODES = {
    '红帅': 1, '红仕': 2, '红相': 3, '红马': 4, '红车': 5, '红炮': 6, '红兵': 7,
    '黑将': -1, '黑士': -2, '黑象': -3, '黑马': -4, '黑车': -5, '黑炮': -6, '黑卒': -7,
    '空': 0
}
PIECE_NAMES = {v: k for k, v in PIECE_CODES.items()}
PIECE_CHARS = {  # 棋子显示字符
    1: '帅', 2: '仕', 3: '相', 4: '马', 5: '车', 6: '炮', 7: '兵',
    -1: '将', -2: '士', -3: '象', -4: '马', -5: '车', -6: '炮', -7: '卒'
}
PIECE_COLORS = {  # 棋子颜色
    1: 'red', 2: 'red', 3: 'red', 4: 'red', 5: 'red', 6: 'red', 7: 'red',
    -1: 'black', -2: 'black', -3: 'black', -4: 'black', -5: 'black', -6: 'black', -7: 'black'
}


# 2. 棋盘核心类（PyTorch存储状态）
class ChineseChessBoard:
    def __init__(self):
        # 用PyTorch张量存储棋盘（10×9）
        self.board = torch.zeros((BOARD_ROWS, BOARD_COLS), dtype=torch.int8)
        self.current_player = 1  # 1=红方（玩家），-1=黑方
        self.selected_pos = None  # 当前选中的棋子坐标(y, x)
        self.valid_moves = []     # 选中棋子的合法移动位置
        self.init_board()

    def init_board(self):
        """初始化中国象棋标准布局"""
        # 红方（第0行：车马相仕帅仕相马车；第2行：兵）
        self.board[0] = torch.tensor([5, 4, 3, 2, 1, 2, 3, 4, 5])
        self.board[2] = torch.tensor([0, 7, 0, 7, 0, 7, 0, 7, 0])
        # 黑方（第9行：车马象士将士象马车；第7行：卒）
        self.board[9] = torch.tensor([-5, -4, -3, -2, -1, -2, -3, -4, -5])
        self.board[7] = torch.tensor([0, -7, 0, -7, 0, -7, 0, -7, 0])

    def get_piece_moves(self, y, x):
        """获取指定棋子的合法移动（完整中国象棋规则）"""
        piece = self.board[y, x].item()
        if piece == 0 or piece * self.current_player <= 0:
            return []  # 空位置/对方棋子无移动权限

        moves = []
        abs_piece = abs(piece)

        # 车：直线移动
        if abs_piece == 5:
            for dx in [1, -1]:
                new_x = x + dx
                while 0 <= new_x < BOARD_COLS:
                    if self.board[y, new_x] == 0:
                        moves.append((y, new_x))
                    else:
                        if self.board[y, new_x] * self.current_player < 0:
                            moves.append((y, new_x))
                        break
                    new_x += dx
            for dy in [1, -1]:
                new_y = y + dy
                while 0 <= new_y < BOARD_ROWS:
                    if self.board[new_y, x] == 0:
                        moves.append((new_y, x))
                    else:
                        if self.board[new_y, x] * self.current_player < 0:
                            moves.append((new_y, x))
                        break
                    new_y += dy

        # 马：日字走法（防蹩马腿）
        elif abs_piece == 4:
            horse_moves = [(y+2,x+1),(y+2,x-1),(y-2,x+1),(y-2,x-1),(y+1,x+2),(y+1,x-2),(y-1,x+2),(y-1,x-2)]
            block_pos = [(y+1,x),(y+1,x),(y-1,x),(y-1,x),(y,x+1),(y,x-1),(y,x+1),(y,x-1)]
            for i, (ny, nx) in enumerate(horse_moves):
                if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS:
                    by, bx = block_pos[i]
                    if self.board[by, bx] == 0 and self.board[ny, nx] * self.current_player <= 0:
                        moves.append((ny, nx))

        # 炮：直线+隔子吃
        elif abs_piece == 6:
            for dx in [1, -1]:
                new_x = x + dx
                has_block = False
                while 0 <= new_x < BOARD_COLS:
                    if self.board[y, new_x] == 0 and not has_block:
                        moves.append((y, new_x))
                    elif self.board[y, new_x] != 0 and not has_block:
                        has_block = True
                    elif self.board[y, new_x] != 0 and has_block:
                        if self.board[y, new_x] * self.current_player < 0:
                            moves.append((y, new_x))
                        break
                    new_x += dx
            for dy in [1, -1]:
                new_y = y + dy
                has_block = False
                while 0 <= new_y < BOARD_ROWS:
                    if self.board[new_y, x] == 0 and not has_block:
                        moves.append((new_y, x))
                    elif self.board[new_y, x] != 0 and not has_block:
                        has_block = True
                    elif self.board[new_y, x] != 0 and has_block:
                        if self.board[new_y, x] * self.current_player < 0:
                            moves.append((new_y, x))
                        break
                    new_y += dy

        # 相/象：田字+不过河
        elif abs_piece == 3:
            if (self.current_player == 1 and y >= 5) or (self.current_player == -1 and y < 5):
                return []
            elephant_moves = [(y+2,x+2),(y+2,x-2),(y-2,x+2),(y-2,x-2)]
            block_pos = [(y+1,x+1),(y+1,x-1),(y-1,x+1),(y-1,x-1)]
            for i, (ny, nx) in enumerate(elephant_moves):
                if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS:
                    by, bx = block_pos[i]
                    if self.board[by, bx] == 0 and self.board[ny, nx] * self.current_player <= 0:
                        moves.append((ny, nx))

        # 仕/士：九宫斜线
        elif abs_piece == 2:
            scholar_moves = [(y+1,x+1),(y+1,x-1),(y-1,x+1),(y-1,x-1)]
            for ny, nx in scholar_moves:
                if (ny in [0,1,2] and nx in [3,4,5]) or (ny in [7,8,9] and nx in [3,4,5]):
                    if self.board[ny, nx] * self.current_player <= 0:
                        moves.append((ny, nx))

        # 帅/将：九宫直线
        elif abs_piece == 1:
            king_moves = [(y+1,x),(y-1,x),(y,x+1),(y,x-1)]
            for ny, nx in king_moves:
                if (ny in [0,1,2] and nx in [3,4,5]) or (ny in [7,8,9] and nx in [3,4,5]):
                    if self.board[ny, nx] * self.current_player <= 0:
                        moves.append((ny, nx))

        # 兵/卒：未过河前直走，过河后可左右
        elif abs_piece == 7:
            if self.current_player == 1:
                if y < 5:
                    if y+1 < BOARD_ROWS and self.board[y+1, x] * self.current_player <= 0:
                        moves.append((y+1, x))
                else:
                    for (dy, dx) in [(1,0),(0,1),(0,-1)]:
                        ny, nx = y+dy, x+dx
                        if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS and self.board[ny, nx] * self.current_player <= 0:
                            moves.append((ny, nx))
            else:
                if y >= 5:
                    if y-1 >= 0 and self.board[y-1, x] * self.current_player <= 0:
                        moves.append((y-1, x))
                else:
                    for (dy, dx) in [(-1,0),(0,1),(0,-1)]:
                        ny, nx = y+dy, x+dx
                        if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS and self.board[ny, nx] * self.current_player <= 0:
                            moves.append((ny, nx))

        return moves

    def move_piece(self, from_y, from_x, to_y, to_x):
        """执行落子（合法则返回True）"""
        if (to_y, to_x) not in self.get_piece_moves(from_y, from_x):
            return False
        self.board[to_y, to_x] = self.board[from_y, from_x]
        self.board[from_y, from_x] = 0
        self.current_player *= -1  # 切换玩家
        self.selected_pos = None
        self.valid_moves = []
        return True


# 3. 可视化与交互逻辑（Matplotlib实现）
class ChessVisualizer:
    def __init__(self, board):
        self.board = board
        self.fig, self.ax = plt.subplots(figsize=(WINDOW_WIDTH/100, WINDOW_HEIGHT/100))
        self.fig.canvas.mpl_connect('button_press_event', self.on_click)
        self.draw_board()
        plt.show()

    def draw_board(self):
        """绘制棋盘、棋子"""
        self.ax.clear()
        self.ax.set_xlim(0, BOARD_COLS * CELL_SIZE)
        self.ax.set_ylim(0, BOARD_ROWS * CELL_SIZE)
        self.ax.set_aspect('equal')
        self.ax.axis('off')

        # 绘制棋盘背景
        for y in range(BOARD_ROWS):
            for x in range(BOARD_COLS):
                # 交替格子颜色（浅棕/深棕）
                color = '#E6C39F' if (y + x) % 2 == 0 else '#D4A76A'
                rect = Rectangle(
                    (x * CELL_SIZE, y * CELL_SIZE), CELL_SIZE, CELL_SIZE,
                    facecolor=color, edgecolor='black', linewidth=1
                )
                self.ax.add_patch(rect)

        # 绘制楚河汉界
        self.ax.text(
            BOARD_COLS * CELL_SIZE / 2, 5 * CELL_SIZE,
            '楚河        汉界', ha='center', va='center',
            fontproperties=font, fontsize=16, color='black'
        )

        # 绘制九宫格斜线
        # 红方九宫（0-2行，3-5列）
        self.ax.plot([3*CELL_SIZE, 5*CELL_SIZE], [0*CELL_SIZE, 2*CELL_SIZE], color='black', linewidth=2)
        self.ax.plot([5*CELL_SIZE, 3*CELL_SIZE], [0*CELL_SIZE, 2*CELL_SIZE], color='black', linewidth=2)
        # 黑方九宫（7-9行，3-5列）
        self.ax.plot([3*CELL_SIZE, 5*CELL_SIZE], [7*CELL_SIZE, 9*CELL_SIZE], color='black', linewidth=2)
        self.ax.plot([5*CELL_SIZE, 3*CELL_SIZE], [7*CELL_SIZE, 9*CELL_SIZE], color='black', linewidth=2)

        # 绘制棋子
        for y in range(BOARD_ROWS):
            for x in range(BOARD_COLS):
                piece = self.board.board[y, x].item()
                if piece != 0:
                    # 棋子圆形背景
                    circle = Circle(
                        (x * CELL_SIZE + CELL_SIZE/2, y * CELL_SIZE + CELL_SIZE/2),
                        CELL_SIZE/2 - 5, facecolor='white', edgecolor=PIECE_COLORS[piece], linewidth=2
                    )
                    self.ax.add_patch(circle)
                    # 棋子文字
                    self.ax.text(
                        x * CELL_SIZE + CELL_SIZE/2, y * CELL_SIZE + CELL_SIZE/2,
                        PIECE_CHARS[piece], ha='center', va='center',
                        fontproperties=font, fontsize=18, color=PIECE_COLORS[piece]
                    )

        # 绘制选中效果（黄色边框）
        if self.board.selected_pos:
            y, x = self.board.selected_pos
            rect = Rectangle(
                (x * CELL_SIZE, y * CELL_SIZE), CELL_SIZE, CELL_SIZE,
                facecolor='none', edgecolor='yellow', linewidth=3
            )
            self.ax.add_patch(rect)
            # 绘制合法移动位置（黄色圆点）
            for (ny, nx) in self.board.valid_moves:
                circle = Circle(
                    (nx * CELL_SIZE + CELL_SIZE/2, ny * CELL_SIZE + CELL_SIZE/2),
                    10, facecolor='yellow', edgecolor='black', linewidth=1
                )
                self.ax.add_patch(circle)

        self.fig.canvas.draw()

    def on_click(self, event):
        """处理鼠标点击事件（选择/移动棋子）"""
        if event.inaxes != self.ax:
            return
        # 转换为棋盘坐标（y是行，x是列）
        x = int(event.xdata // CELL_SIZE)
        y = int(event.ydata // CELL_SIZE)
        if not (0 <= x < BOARD_COLS and 0 <= y < BOARD_ROWS):
            return

        # 选择自己的棋子
        piece = self.board.board[y, x].item()
        if piece * self.board.current_player > 0:
            self.board.selected_pos = (y, x)
            self.board.valid_moves = self.board.get_piece_moves(y, x)
        # 移动选中的棋子
        elif self.board.selected_pos:
            from_y, from_x = self.board.selected_pos
            if (y, x) in self.board.valid_moves:
                self.board.move_piece(from_y, from_x, y, x)
        else:
            self.board.selected_pos = None
            self.board.valid_moves = []

        self.draw_board()


# 启动游戏
if __name__ == "__main__":
    chess_board = ChineseChessBoard()
    visualizer = ChessVisualizer(chess_board)
